import numpy
import pprint
from matplotlib import pyplot
dataset = numpy.loadtxt('./exp2.txt', delimiter='\t', encoding='utf-8')

def linearSolve(data):#这里使用普通的线性拟合
    shape=data.shape
    Y=data[:,-1]
    X=numpy.ones(shape)
    X[:,1:]=data[:,0:-1]
    xTx=X.T@X
    if numpy.linalg.det(xTx)==0:
        print('this matrix is singular')
    W = (numpy.linalg.inv((xTx)) @ X.T) @ Y
    return W,X,Y
def Err(data):
    W,X,Y=linearSolve(data)
    Y_pre=X@W
    return numpy.sum((Y-Y_pre)**2)
def leaf(data):
    W, X, Y = linearSolve(data)
    return W
#模型树的划分依据为进行拟合后的误差,可以根据需要自行选择合理的拟合方式
def get_best_split(data,slimit,nlimit):
    if numpy.unique(data[:, -1]).shape[0] == 1:
        return None, leaf(data)  # 如果值相同则不用划分，直接返回值
    shape = data.shape
    S = Err(data)  # 计算数据原来的总方差
    bestS = numpy.inf
    spIndex = 0
    spValue = 0
    for i in range(shape[1] - 1):
        for j in range(shape[0]):
            right = data[data[:, i] > data[j, i]]
            left = data[data[:, i] <= data[j, i]]
            if right.shape[0]<=1 or left.shape[0]<=1:#如果只有一个数据，则不参与划分，因为一个数据无法进行拟合
                continue
            newS = Err(right)+Err(left)
            if newS < bestS:
                bestS = newS
                spIndex = i
                spValue = data[j, i]
    # 划分后的方差变化太小则取消此次划分
    if S - bestS < slimit:
        return None, leaf(data)
    # 划分后的数据集数量太小则取消此次划分
    right = data[data[:, spIndex] > spValue]
    left = data[data[:, spIndex] <= spValue]#用最好的划分方式进行划分
    if right.shape[0] < nlimit or left.shape[0] < nlimit:
        return None, leaf(data)
    return spIndex, spValue
def model_tree(data,slimit,nlimit):
    spIndex, spValue = get_best_split(data,slimit,nlimit)
    if spIndex is None:
        return spValue
    tree = {'spIndex': spIndex, 'spValue': spValue}
    right = data[data[:, spIndex] > spValue]
    left = data[data[:, spIndex] <= spValue]
    tree['right'] = model_tree(right,slimit,nlimit)
    tree['left'] =model_tree(left,slimit,nlimit)
    return tree
#对数据进行预测
def predict(tree,X):
    if not isinstance(tree,dict):
        return tree@numpy.array([1,X[0]])
    if X[tree['spIndex']]<=tree['spValue']:
        return predict(tree['left'],X)
    else:
        return predict(tree['right'],X)
tree=model_tree(dataset,1,10)
print(tree)
li=[]
x=numpy.arange(0,1.1,0.1)
for i in x:
    li.append(predict(tree,numpy.array([i])))
pyplot.scatter(dataset[:, 0], dataset[:, 1],s=1)
pyplot.plot(x,li)
pyplot.show()